{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "narrative-sailing",
   "metadata": {
    "papermill": {
     "duration": 0.040767,
     "end_time": "2021-03-24T14:13:51.082088",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.041321",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Inference on Predictive and Causal Effects in High-Dimensional Nonlinear Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ready-appearance",
   "metadata": {
    "papermill": {
     "duration": 0.036878,
     "end_time": "2021-03-24T14:13:51.155414",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.118536",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Impact of 401(k) on  Financial Wealth\n",
    "\n",
    "As a practical illustration of the methods developed in this lecture, we consider estimation of the effect of 401(k) eligibility and participation \n",
    "on accumulated assets. 401(k) plans are pension accounts sponsored by employers. The key problem in determining the effect of participation in 401(k) plans on accumulated assets is saver heterogeneity coupled with the fact that the decision to enroll in a 401(k) is non-random. It is generally recognized that some people have a higher preference for saving than others. It also seems likely that those individuals with high unobserved preference for saving would be most likely to choose to participate in tax-advantaged retirement savings plans and would tend to have otherwise high amounts of accumulated assets. The presence of unobserved savings preferences with these properties then implies that conventional estimates that do not account for saver heterogeneity and endogeneity of participation will be biased upward, tending to overstate the savings effects of 401(k) participation.\n",
    "\n",
    "One can argue that eligibility for enrolling in a 401(k) plan in this data can be taken as exogenous after conditioning on a few observables of which the most important for their argument is income. The basic idea is that, at least around the time 401(k)’s initially became available, people were unlikely to be basing their employment decisions on whether an employer offered a 401(k) but would instead focus on income and other aspects of the job. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "divine-phoenix",
   "metadata": {
    "papermill": {
     "duration": 0.036497,
     "end_time": "2021-03-24T14:13:51.229370",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.192873",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Data\n",
    "\n",
    "The data set can be downloaded from the github repo\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75662e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import relevant packages\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.model_selection import cross_val_score, cross_val_predict, KFold\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.linear_model import LassoCV, RidgeCV, ElasticNetCV, LinearRegression, Ridge, Lasso, LogisticRegressionCV\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier\n",
    "from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier\n",
    "from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier\n",
    "import patsy\n",
    "import warnings\n",
    "from sklearn.base import BaseEstimator, clone\n",
    "import statsmodels.api as sm\n",
    "from IPython.display import Markdown\n",
    "import wget\n",
    "import os\n",
    "import seaborn as sns\n",
    "warnings.simplefilter('ignore')\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e66e77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = \"https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/401k.csv\"\n",
    "data = pd.read_csv(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a2d086",
   "metadata": {},
   "outputs": [],
   "source": [
    "data.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00884061",
   "metadata": {},
   "outputs": [],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3dc85c",
   "metadata": {},
   "outputs": [],
   "source": [
    "readme = \"https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/401k.md\"\n",
    "filename = wget.download(readme)\n",
    "display(Markdown(open(filename, 'r').read()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "looking-invention",
   "metadata": {
    "papermill": {
     "duration": 0.037509,
     "end_time": "2021-03-24T14:13:52.250583",
     "exception": false,
     "start_time": "2021-03-24T14:13:52.213074",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The data consist of 9,915 observations at the household level drawn from the 1991 Survey of Income and Program Participation (SIPP).  All the variables are referred to 1990. We use net financial assets (*net\\_tfa*) as the outcome variable, $Y$,  in our analysis. The net financial assets are computed as the sum of IRA balances, 401(k) balances, checking accounts, saving bonds, other interest-earning accounts, other interest-earning assets, stocks, and mutual funds less non mortgage debts. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "received-nutrition",
   "metadata": {
    "papermill": {
     "duration": 0.038843,
     "end_time": "2021-03-24T14:13:52.327055",
     "exception": false,
     "start_time": "2021-03-24T14:13:52.288212",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $9915$ individuals, $3682$ are eligible to participate in the program. The variable *e401* indicates eligibility and *p401* indicates participation, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d26d552",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.countplot(data['e401'],)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "material-sending",
   "metadata": {
    "papermill": {
     "duration": 0.039406,
     "end_time": "2021-03-24T14:13:53.435306",
     "exception": false,
     "start_time": "2021-03-24T14:13:53.395900",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Eligibility is highly associated with financial wealth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7741bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.displot(data=data, x='net_tfa', kind='kde', col='e401', hue='e401', fill=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "awful-antigua",
   "metadata": {
    "papermill": {
     "duration": 0.040645,
     "end_time": "2021-03-24T14:13:54.245822",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.205177",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The unconditional APE of e401 is about $19559$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef070d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1 = data[data['e401'] == 1]['net_tfa']\n",
    "e0 = data[data['e401'] == 0]['net_tfa']\n",
    "print(f'{np.mean(e1) - np.mean(e0):.0f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cross-priority",
   "metadata": {
    "papermill": {
     "duration": 0.041187,
     "end_time": "2021-03-24T14:13:54.437695",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.396508",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $3682$ individuals that  are eligible, $2594$ decided to participate in the program. The unconditional APE of p401 is about $27372$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33cd014e",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1 = data[data['p401'] == 1]['net_tfa']\n",
    "e0 = data[data['p401'] == 0]['net_tfa']\n",
    "print(f'{np.mean(e1) - np.mean(e0):.0f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suitable-vulnerability",
   "metadata": {
    "papermill": {
     "duration": 0.041678,
     "end_time": "2021-03-24T14:13:54.640532",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.598854",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As discussed, these estimates are biased since they do not account for saver heterogeneity and endogeneity of participation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9354414",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = data['net_tfa'].values\n",
    "D = data['e401'].values\n",
    "D2 = data['p401'].values\n",
    "D3 = data['a401'].values\n",
    "X = data.drop(['e401', 'p401', 'a401', 'tw', 'tfa', 'net_tfa', 'tfa_he',\n",
    "               'hval', 'hmort', 'hequity',\n",
    "               'nifa', 'net_nifa', 'net_n401', 'ira',\n",
    "               'dum91', 'icat', 'ecat', 'zhat',\n",
    "               'i1', 'i2', 'i3', 'i4', 'i5', 'i6', 'i7',\n",
    "               'a1', 'a2', 'a3', 'a4', 'a5'], axis=1)\n",
    "X.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c3e489",
   "metadata": {},
   "source": [
    "### We define a transformer that constructs the engineered features for controls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2cfdf4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install formulaic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b1283d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.base import TransformerMixin, BaseEstimator\n",
    "from formulaic import Formula\n",
    "\n",
    "class FormulaTransformer(TransformerMixin, BaseEstimator):\n",
    "    \n",
    "    def __init__(self, formula, array=False):\n",
    "        self.formula = formula\n",
    "        self.array = array\n",
    "    \n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "    \n",
    "    def transform(self, X, y=None):\n",
    "        df = Formula(self.formula).get_model_matrix(X)\n",
    "        if self.array:\n",
    "            return df.values\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824ee320",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = FormulaTransformer(\"0 + poly(age, degree=6, raw=True) + poly(inc, degree=8, raw=True) \"\n",
    "                                 \"+ poly(educ, degree=4, raw=True) + poly(fsize, degree=2, raw=True) \"\n",
    "                                 \"+ male + marr + twoearn + db + pira + hown\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b35bbd34",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer.fit_transform(X).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c5e20e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = FormulaTransformer(\"0 + poly(age, degree=6, raw=True) + poly(inc, degree=8, raw=True) \"\n",
    "                                 \"+ poly(educ, degree=4, raw=True) + poly(fsize, degree=2, raw=True) \"\n",
    "                                 \"+ marr + twoearn + db + pira + hown\", array=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a745a6f2",
   "metadata": {},
   "source": [
    "## Estimating the ATE of 401(k) Eligibility on Net Financial Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10538ad3",
   "metadata": {},
   "source": [
    "We are interested in valid estimators of the average treatment effect of `e401` and `p401` on `net_tfa`. We start using ML approaches to estimate the function $g_0$ and $m_0$ in the following PLR model:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fcf3165",
   "metadata": {},
   "source": [
    "\\begin{eqnarray}\n",
    " &  Y = D\\theta_0 + g_0(X) + \\zeta,  &  E[\\zeta \\mid D,X]= 0,\\\\\n",
    " & D = m_0(X) +  V,   &  E[V \\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42354b51",
   "metadata": {},
   "source": [
    "### Double Lasso without Sample Splitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e3c38b",
   "metadata": {},
   "source": [
    "We first look at the treatment effect of e401 on net total financial assets. We give estimates of the ATE and ATT that corresponds to the linear model\n",
    "\n",
    "\\begin{equation*}\n",
    "Y = D \\alpha + f(X)'\\beta+ \\epsilon,\n",
    "\\end{equation*}\n",
    "\n",
    "where $f(X)$ includes indicators of marital status, two-earner status, defined benefit pension status, IRA participation status, and home ownership status, and  orthogonal polynomials of degrees 2, 4, 6 and 8 in family size, education, age and  income, respectively. The dimensions of $f(X)$ is 25. \n",
    "\n",
    "In the first step, we report estimates of the average treatment effect (ATE) of 401(k) eligibility on net financial assets both in the partially linear regression (PLR) model and in the interactive regression model (IRM) allowing for heterogeneous treatment effects. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "051af0ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "modely = make_pipeline(transformer, StandardScaler(), LassoCV())\n",
    "modeld = make_pipeline(transformer, StandardScaler(), LassoCV())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53bc1c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "resy = y - modely.fit(X, y).predict(X)\n",
    "resD = D - modeld.fit(X, D).predict(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51244ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean(resy * resD) / np.mean(resD**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d198da92",
   "metadata": {},
   "source": [
    "### Double ML in PLR with Cross-Fitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "045487a5",
   "metadata": {},
   "source": [
    "We define a simple dml function that is parameterized by arbitrary ML models and returns the treatment effect and other useful quantities of the analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee13ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dml(X, D, y, modely, modeld, *, nfolds, classifier=False):\n",
    "    '''\n",
    "    DML for the Partially Linear Model setting with cross-fitting\n",
    "    \n",
    "    Input\n",
    "    -----\n",
    "    X: the controls\n",
    "    D: the treatment\n",
    "    y: the outcome\n",
    "    modely: the ML model for predicting the outcome y\n",
    "    modeld: the ML model for predicting the treatment D\n",
    "    nfolds: the number of folds in cross-fitting\n",
    "    classifier: bool, whether the modeld is a classifier or a regressor\n",
    "    \n",
    "    Output\n",
    "    ------\n",
    "    point: the point estimate of the treatment effect of D on y\n",
    "    stderr: the standard error of the treatment effect\n",
    "    yhat: the cross-fitted predictions for the outcome y\n",
    "    Dhat: the cross-fitted predictions for the treatment D\n",
    "    resy: the outcome residuals\n",
    "    resD: the treatment residuals\n",
    "    epsilon: the final residual-on-residual OLS regression residual\n",
    "    '''\n",
    "    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123) # shuffled k-folds\n",
    "    yhat = cross_val_predict(modely, X, y, cv=cv, n_jobs=-1) # out-of-fold predictions for y\n",
    "    # out-of-fold predictions for D\n",
    "    # use predict or predict_proba dependent on classifier or regressor for D\n",
    "    if classifier: \n",
    "        Dhat = cross_val_predict(modeld, X, D, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n",
    "    else:\n",
    "        Dhat = cross_val_predict(modeld, X, D, cv=cv, n_jobs=-1)\n",
    "    # calculate outcome and treatment residuals\n",
    "    resy = y - yhat\n",
    "    resD = D - Dhat\n",
    "    # final stage ols based point estimate and standard error\n",
    "    point = np.mean(resy * resD) / np.mean(resD**2)\n",
    "    epsilon = resy - point * resD\n",
    "    var = np.mean(epsilon**2 * resD**2) / np.mean(resD**2)**2\n",
    "    stderr = np.sqrt(var / X.shape[0])\n",
    "    return point, stderr, yhat, Dhat, resy, resD, epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35c70f74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summary(point, stderr, yhat, Dhat, resy, resD, epsilon, X, D, y, *, name):\n",
    "    '''\n",
    "    Convenience summary function that takes the results of the DML function\n",
    "    and summarizes several estimation quantities and performance metrics.\n",
    "    '''\n",
    "    return pd.DataFrame({'estimate': point, # point estimate\n",
    "                         'stderr': stderr, # standard error\n",
    "                         'lower': point - 1.96*stderr, # lower end of 95% confidence interval\n",
    "                         'upper': point + 1.96*stderr, # upper end of 95% confidence interval\n",
    "                         'rmse y': np.sqrt(np.mean(resy**2)), # RMSE of model that predicts outcome y\n",
    "                         'rmse D': np.sqrt(np.mean(resD**2)), # RMSE of model that predicts treatment D\n",
    "                         'accuracy D': np.mean(np.abs(resD) < .5), # binary classification accuracy of model for D\n",
    "                         }, index=[name])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0e6276a",
   "metadata": {},
   "source": [
    "#### Double Lasso with Cross-Fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f57b4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cv = KFold(n_splits=5, shuffle=True, random_state=123)\n",
    "lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n",
    "lassod = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n",
    "result = dml(X, D, y, lassoy, lassod, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "265fb305",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = summary(*result, X, D, y, name='double lasso')\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30cb77e8",
   "metadata": {},
   "source": [
    "#### Using a Penalized Logistic Regression for D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba3ba66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cv = KFold(n_splits=5, shuffle=True, random_state=123)\n",
    "lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n",
    "lgrd = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n",
    "result = dml(X, D, y, lassoy, lgrd, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbaa344",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='lasso/logistic'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-gregory",
   "metadata": {
    "papermill": {
     "duration": 0.048205,
     "end_time": "2021-03-24T14:17:09.946894",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.898689",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Then, we repeat this procedure for various machine learning methods."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4791084",
   "metadata": {},
   "source": [
    "### Random Forests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0d4030",
   "metadata": {},
   "outputs": [],
   "source": [
    "rfy = make_pipeline(transformer, RandomForestRegressor(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n",
    "rfd = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n",
    "result = dml(X, D, y, rfy, rfd, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5073e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='random forest'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1adab609",
   "metadata": {},
   "source": [
    "### Decision Trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a3266c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtry = make_pipeline(transformer, DecisionTreeRegressor(min_samples_leaf=10, ccp_alpha=.001))\n",
    "dtrd = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n",
    "result = dml(X, D, y, dtry, dtrd, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e82ddc9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='decision tree'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "415d5b98",
   "metadata": {},
   "source": [
    "### Boosted Trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9016b1a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbfy = make_pipeline(transformer, GradientBoostingRegressor(max_depth=2, n_iter_no_change=5))\n",
    "gbfd = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n",
    "result = dml(X, D, y, gbfy, gbfd, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10f4e992",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='boosted forest'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identical-smith",
   "metadata": {
    "papermill": {
     "duration": 0.056641,
     "end_time": "2021-03-24T14:18:16.044675",
     "exception": false,
     "start_time": "2021-03-24T14:18:15.988034",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The best model with lowest RMSE in both equation is the PLR model estimated via lasso."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c8dc0c4",
   "metadata": {},
   "source": [
    "### AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d0d2b4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flaml import AutoML\n",
    "flamly = make_pipeline(transformer, AutoML(time_budget=100, task='regression', early_stop=True,\n",
    "                                    eval_method='cv', n_splits=3, metric='r2', verbose=3))\n",
    "flamld = make_pipeline(transformer, AutoML(time_budget=100, task='classification', early_stop=True,\n",
    "                                           eval_method='cv', n_splits=3, metric='r2', verbose=3))\n",
    "result = dml(X, D, y, flamly, flamld, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2ba772",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='automl'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bc171ce",
   "metadata": {},
   "source": [
    "# Semi-Cross-Fitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9d4e9e",
   "metadata": {},
   "source": [
    "To avoid the computational cost of performing model selection within each fold, assuming that we don't select among an exponential set of hyperparameters/models (in the number of samples); it is ok to perform model selection using all the data and then perform cross-fitting with the selected model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f64eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamly = make_pipeline(transformer, AutoML(time_budget=100, task='regression', early_stop=True,\n",
    "                                    eval_method='cv', n_splits=3, metric='r2', verbose=0))\n",
    "flamld = make_pipeline(transformer, AutoML(time_budget=100, task='classification', early_stop=True,\n",
    "                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cffa502",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamly.fit(X, y)\n",
    "besty = make_pipeline(transformer, clone(flamly[-1].best_model_for_estimator(flamly[-1].best_estimator)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c2a08a",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamld.fit(X, D)\n",
    "bestd = make_pipeline(transformer, clone(flamld[-1].best_model_for_estimator(flamld[-1].best_estimator)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5a1b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = dml(X, D, y, besty, bestd, nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9855e83e",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result,  X, D, y, name='automl (semi-cfit)'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b97cb96c",
   "metadata": {},
   "source": [
    "### Semi-Crossfitting with Stacking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5214c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dml_dirty(X, D, y, modely_list, modeld_list, *,\n",
    "              stacker=LinearRegression(), nfolds, classifier=False):\n",
    "    '''\n",
    "    DML for the Partially Linear Model setting with semi-cross-fitting\n",
    "    \n",
    "    Input\n",
    "    -----\n",
    "    X: the controls\n",
    "    D: the treatment\n",
    "    y: the outcome\n",
    "    modely: the ML model for predicting the outcome y\n",
    "    modeld: the ML model for predicting the treatment D\n",
    "    stacker: model used to aggregate predictions of each of the base models\n",
    "    nfolds: the number of folds in cross-fitting\n",
    "    classifier: bool, whether the modeld is a classifier or a regressor\n",
    "    \n",
    "    Output\n",
    "    ------\n",
    "    point: the point estimate of the treatment effect of D on y\n",
    "    stderr: the standard error of the treatment effect\n",
    "    yhat: the cross-fitted predictions for the outcome y\n",
    "    Dhat: the cross-fitted predictions for the treatment D\n",
    "    resy: the outcome residuals\n",
    "    resD: the treatment residuals\n",
    "    epsilon: the final residual-on-residual OLS regression residual\n",
    "    '''\n",
    "    # construct out-of-fold predictions for each model\n",
    "    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123)\n",
    "    yhats = np.array([cross_val_predict(modely, X, y, cv=cv, n_jobs=-1) for modely in modely_list]).T\n",
    "    if classifier:\n",
    "        Dhats = np.array([cross_val_predict(modeld, X, D, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n",
    "                         for modeld in modeld_list]).T\n",
    "    else:\n",
    "        Dhats = np.array([cross_val_predict(modeld, X, D, cv=cv, n_jobs=-1) for modeld in modeld_list]).T\n",
    "    # calculate stacked residuals by finding optimal coefficients\n",
    "    # and weigthing out-of-sample predictions by these coefficients\n",
    "    yhat = stacker.fit(yhats, y).predict(yhats)\n",
    "    Dhat = stacker.fit(Dhats, D).predict(Dhats)\n",
    "    resy = y - yhat\n",
    "    resD = D - Dhat\n",
    "    # go with the stacked residuals\n",
    "    point = np.mean(resy * resD) / np.mean(resD**2)\n",
    "    epsilon = resy - point * resD\n",
    "    var = np.mean(epsilon**2 * resD**2) / np.mean(resD**2)**2\n",
    "    stderr = np.sqrt(var / X.shape[0])\n",
    "    return point, stderr, yhat, Dhat, resy, resD, epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a671f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = dml_dirty(X, D, y, [lassoy, rfy, dtry, gbfy], [lgrd, rfd, dtrd, gbfd],\n",
    "                   nfolds=3, classifier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6bfa361",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = table.append(summary(*result, X, D, y, name='stacked (semi-cfit)'))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comprehensive-graphics",
   "metadata": {
    "papermill": {
     "duration": 0.056567,
     "end_time": "2021-03-24T14:18:16.296870",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.240303",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Interactive Regression Model (IRM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "square-craps",
   "metadata": {
    "papermill": {
     "duration": 0.057315,
     "end_time": "2021-03-24T14:18:16.410105",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.352790",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Next, we consider estimation of average treatment effects when treatment effects are fully heterogeneous:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "large-welcome",
   "metadata": {
    "papermill": {
     "duration": 0.05644,
     "end_time": "2021-03-24T14:18:16.525113",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.468673",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    " \\begin{eqnarray}\\label{eq: HetPL1}\n",
    " & Y  = g_0(D, X) + U,  &  \\quad E[U \\mid X, D]= 0,\\\\\n",
    "  & D  = m_0(X) + V,  & \\quad  E[V\\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "perfect-reliance",
   "metadata": {
    "papermill": {
     "duration": 0.056073,
     "end_time": "2021-03-24T14:18:16.637385",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.581312",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "To reduce the disproportionate impact of extreme propensity score weights in the interactive model\n",
    "we trim the propensity scores which are close to the bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a0b5d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dr(X, D, y, modely0, modely1, modeld, *, trimming=0.01, nfolds):\n",
    "    '''\n",
    "    DML for the Interactive Regression Model setting (Doubly Robust Learning)\n",
    "    with cross-fitting\n",
    "    \n",
    "    Input\n",
    "    -----\n",
    "    X: the controls\n",
    "    D: the treatment\n",
    "    y: the outcome\n",
    "    modely0: the ML model for predicting the outcome y in the control population\n",
    "    modely1: the ML model for predicting the outcome y in the treated population\n",
    "    modeld: the ML model for predicting the treatment D\n",
    "    trimming: threshold below which to trim propensities\n",
    "    nfolds: the number of folds in cross-fitting\n",
    "    \n",
    "    Output\n",
    "    ------\n",
    "    point: the point estimate of the treatment effect of D on y\n",
    "    stderr: the standard error of the treatment effect\n",
    "    yhat: the cross-fitted predictions for the outcome y\n",
    "    Dhat: the cross-fitted predictions for the outcome D\n",
    "    resy: the outcome residuals\n",
    "    resD: the treatment residuals\n",
    "    drhat: the doubly robust quantity for each sample\n",
    "    '''\n",
    "    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123)\n",
    "    yhat0, yhat1 = np.zeros(y.shape), np.zeros(y.shape)\n",
    "    # we will fit a model E[Y| D, X] by fitting a separate model for D==0\n",
    "    # and a separate model for D==1.\n",
    "    for train, test in cv.split(X, y):\n",
    "        # train a model on training data that received treatment zero and predict on all data in test set\n",
    "        yhat0[test] = clone(modely0).fit(X.iloc[train][D[train]==0], y[train][D[train]==0]).predict(X.iloc[test])\n",
    "        # train a model on training data that received treatment one and predict on all data in test set\n",
    "        yhat1[test] = clone(modely1).fit(X.iloc[train][D[train]==1], y[train][D[train]==1]).predict(X.iloc[test])\n",
    "    # prediction for observed treatment\n",
    "    yhat = yhat0 * (1 - D) + yhat1 * D\n",
    "    # propensity scores\n",
    "    Dhat = cross_val_predict(modeld, X, D, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n",
    "    Dhat = np.clip(Dhat, trimming, 1 - trimming)\n",
    "    # doubly robust quantity for every sample\n",
    "    drhat = yhat1 - yhat0 + (y - yhat) * (D/Dhat - (1 - D)/(1 - Dhat))\n",
    "    point = np.mean(drhat)\n",
    "    var = np.var(drhat)\n",
    "    stderr = np.sqrt(var / X.shape[0])\n",
    "    return point, stderr, yhat, Dhat, y - yhat, D - Dhat, drhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32a89903",
   "metadata": {},
   "outputs": [],
   "source": [
    "cv = KFold(n_splits=5, shuffle=True, random_state=123)\n",
    "lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n",
    "lgrd = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n",
    "result = dr(X, D, y, lassoy, lassoy, lgrd, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8fa94fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = summary(*result, X, D, y, name='lasso/logistic')\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b0fe147",
   "metadata": {},
   "outputs": [],
   "source": [
    "rfy = make_pipeline(transformer, RandomForestRegressor(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n",
    "rfd = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n",
    "result = dr(X, D, y, rfy, rfy, rfd, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "259efb1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = tabledr.append(summary(*result,  X, D, y, name='random forest'))\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5335633c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtry = make_pipeline(transformer, DecisionTreeRegressor(min_samples_leaf=10, ccp_alpha=.001))\n",
    "dtrd = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n",
    "result = dr(X, D, y, dtry, dtry, dtrd, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c5d5a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = tabledr.append(summary(*result,  X, D, y, name='decision tree'))\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f9338c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbfy = make_pipeline(transformer, GradientBoostingRegressor(max_depth=2, n_iter_no_change=5))\n",
    "gbfd = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n",
    "result = dr(X, D, y, gbfy, gbfy, gbfd, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d275ed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = tabledr.append(summary(*result, X, D, y, name='boosted forest'))\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a6e0994",
   "metadata": {},
   "source": [
    "These estimates that flexibly account for confounding are substantially attenuated relative to the baseline estimate (19559) that does not account for confounding. They suggest much smaller causal effects of 401(k) eligiblity on financial asset holdings."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8eadf06",
   "metadata": {},
   "source": [
    "# Semi-Cross-Fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fee26e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flaml import AutoML\n",
    "flamly0 = make_pipeline(transformer, AutoML(time_budget=60, task='regression', early_stop=True,\n",
    "                                     eval_method='cv', n_splits=3, metric='r2', verbose=0))\n",
    "flamly1 = make_pipeline(transformer, AutoML(time_budget=60, task='regression', early_stop=True,\n",
    "                                     eval_method='cv', n_splits=3, metric='r2', verbose=0))\n",
    "flamld = make_pipeline(transformer, AutoML(time_budget=60, task='classification', early_stop=True,\n",
    "                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22518880",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamly0.fit(X[D==0], y[D==0])\n",
    "besty0 = make_pipeline(transformer, clone(flamly0[-1].best_model_for_estimator(flamly0[-1].best_estimator)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd04a3cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamly1.fit(X[D==1], y[D==1])\n",
    "besty1 = make_pipeline(transformer, clone(flamly1[-1].best_model_for_estimator(flamly1[-1].best_estimator)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5088756e",
   "metadata": {},
   "outputs": [],
   "source": [
    "flamld.fit(X, D)\n",
    "bestd = make_pipeline(transformer, clone(flamld[-1].best_model_for_estimator(flamld[-1].best_estimator)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a415bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = dr(X, D, y, besty0, besty1, bestd, nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30800bbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = tabledr.append(summary(*result,  X, D, y, name='automl (semi-cfit)'))\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9325941a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dr_dirty(X, D, y, modely0_list, modely1_list, modeld_list, *,\n",
    "             stacker=LinearRegression(), trimming=0.01, nfolds):\n",
    "    '''\n",
    "    DML for the Interactive Regression Model setting (Doubly Robust Learning)\n",
    "    with cross-fitting\n",
    "    \n",
    "    Input\n",
    "    -----\n",
    "    X: the controls\n",
    "    D: the treatment\n",
    "    y: the outcome\n",
    "    modely_list: list of ML models for predicting the outcome y\n",
    "    modeld_list: list of ML models for predicting the treatment D\n",
    "    stacker: model used to aggregate predictions of each of the base models\n",
    "    trimming: threshold below which to trim propensities\n",
    "    nfolds: the number of folds in cross-fitting\n",
    "    \n",
    "    Output\n",
    "    ------\n",
    "    point: the point estimate of the treatment effect of D on y\n",
    "    stderr: the standard error of the treatment effect\n",
    "    yhat: the cross-fitted predictions for the outcome y\n",
    "    Dhat: the cross-fitted predictions for the outcome D\n",
    "    resy: the outcome residuals\n",
    "    resD: the treatment residuals\n",
    "    drhat: the doubly robust quantity for each sample\n",
    "    '''\n",
    "    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123)\n",
    "    \n",
    "    # we will fit a model E[Y| D, X] by fitting a separate model for D==0\n",
    "    # and a separate model for D==1. We do that for each model type in modely_list\n",
    "    yhats0, yhats1 = np.zeros((y.shape[0], len(modely0_list))), np.zeros((y.shape[0], len(modely1_list)))\n",
    "    for train, test in cv.split(X, y):\n",
    "        for it, modely0 in enumerate(modely0_list):\n",
    "            yhats0[test, it] = clone(modely0).fit(X.iloc[train][D[train]==0], y[train][D[train]==0]).predict(X.iloc[test])\n",
    "        for it, modely1 in enumerate(modely1_list):\n",
    "            yhats1[test, it] = clone(modely1).fit(X.iloc[train][D[train]==1], y[train][D[train]==1]).predict(X.iloc[test])\n",
    "    \n",
    "    # calculate stacking weights for the outcome model for each population\n",
    "    # and combine the outcome model predictions\n",
    "    yhat0 = clone(stacker).fit(yhats0[D==0], y[D==0]).predict(yhats0)\n",
    "    yhat1 = clone(stacker).fit(yhats1[D==1], y[D==1]).predict(yhats1)\n",
    "    \n",
    "    # prediction for observed treatment using the stacked model\n",
    "    yhat = yhat0 * (1 - D) + yhat1 * D\n",
    "\n",
    "    # propensity scores\n",
    "    Dhats = np.array([cross_val_predict(modeld, X, D, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n",
    "                     for modeld in modeld_list]).T\n",
    "    # construct coefficients on each model based on stacker\n",
    "    Dhat = clone(stacker).fit(Dhats, D).predict(Dhats)\n",
    "    # trim propensities\n",
    "    Dhat = np.clip(Dhat, trimming, 1 - trimming)\n",
    "\n",
    "    # doubly robust quantity for every sample\n",
    "    drhat = yhat1 - yhat0 + (y - yhat) * (D/Dhat - (1 - D)/(1 - Dhat))\n",
    "    point = np.mean(drhat)\n",
    "    var = np.var(drhat)\n",
    "    stderr = np.sqrt(var / X.shape[0])\n",
    "    return point, stderr, yhat, Dhat, y - yhat, D - Dhat, drhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5401f09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = dr_dirty(X, D, y, [lassoy, rfy, dtry, gbfy], [lassoy, rfy, dtry, gbfy], [lgrd, rfd, dtrd, gbfd], nfolds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1606f1bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr = tabledr.append(summary(*result,  X, D, y, name='stacked (semi-cfit)'))\n",
    "tabledr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743cde38",
   "metadata": {},
   "source": [
    "We can compare the results between the PLR model and IRM model. We find that the effect under the IRM model is typically of lower value, potentially implying that the PLR model leaves some degree of confounding due to mis-specification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce30440",
   "metadata": {},
   "outputs": [],
   "source": [
    "table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fcee1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tabledr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af5867dc",
   "metadata": {},
   "source": [
    "# Using the EconML Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "401bc6b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install econml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebd88f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for these libraries we will just pre-featurize the controls\n",
    "W = StandardScaler().fit_transform(transformer.fit_transform(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3ba9f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from econml.dml import LinearDML\n",
    "cv = KFold(n_splits=5, shuffle=True, random_state=123)\n",
    "ldml = LinearDML(model_y=LassoCV(cv=cv), model_t=LogisticRegressionCV(cv=cv),\n",
    "                 cv=3, discrete_treatment=True, random_state=123).fit(y, D, W=W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a52d86c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ldml.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d6c5b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# r2scores\n",
    "r2scorey = np.mean(ldml.nuisance_scores_y)\n",
    "r2scored = np.mean(ldml.nuisance_scores_t)\n",
    "r2scorey, r2scored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "473f350b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from econml.dr import LinearDRLearner\n",
    "\n",
    "# dr learner in econml fits a single regression function of Y from X, D\n",
    "# using all the data\n",
    "dr = LinearDRLearner(model_regression=LassoCV(cv=cv),\n",
    "                     model_propensity=LogisticRegressionCV(cv=cv),\n",
    "                     cv=3, min_propensity=0.01, random_state=123).fit(y, D, W=W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea53e99a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dr.summary(T=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519426af",
   "metadata": {},
   "outputs": [],
   "source": [
    "from econml.dr import LinearDRLearner\n",
    "from econml.utilities import SeparateModel\n",
    "\n",
    "# to implement the separate regression models we need use the separate model wrapper\n",
    "# that splits the data based on the last covariate (in this case D) and fits a separate\n",
    "# model for each group. The input to the wrapper is the model to use for each value\n",
    "# of the last covariate.\n",
    "dr = LinearDRLearner(model_regression=SeparateModel(LassoCV(cv=cv), LassoCV(cv=cv)),\n",
    "                     model_propensity=LogisticRegressionCV(cv=cv),\n",
    "                     cv=3, min_propensity=0.01, random_state=123).fit(y, D, W=W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3794606",
   "metadata": {},
   "outputs": [],
   "source": [
    "dr.summary(T=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ca94a11",
   "metadata": {},
   "source": [
    "# Using the DoubleML library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea90e4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install doubleml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9930faf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from doubleml import DoubleMLData\n",
    "dml_data = DoubleMLData.from_arrays(W, y, D)\n",
    "print(dml_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c146dd15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import doubleml as dml\n",
    "\n",
    "dml_plr_obj = dml.DoubleMLPLR(dml_data, LassoCV(cv=cv), LogisticRegressionCV(cv=cv), n_folds=3)\n",
    "print(dml_plr_obj.fit())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e7b7b7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_irm_obj = dml.DoubleMLIRM(dml_data, LassoCV(cv=cv), LogisticRegressionCV(cv=cv), n_folds=3)\n",
    "\n",
    "print(dml_irm_obj.fit())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd7a8e80",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 537.547458,
   "end_time": "2021-03-24T14:22:44.931595",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2021-03-24T14:13:47.384137",
   "version": "2.3.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
